from typing import Union, Dict, Any
from fastapi import APIRouter, HTTPException, BackgroundTasks
from app.services.mvhand_api_service import MVHANDService
from app.schemas.dhgam.mvhand_base_schema import AlgorithmRequest, AlgorithmResponse, AlgorithmMiddleResponse, TaskStatus
from app.tools.redis_client import redis_client
from app.utils.logger import logger

router = APIRouter()
service = MVHANDService()

@router.post("/create", response_model=AlgorithmMiddleResponse)
async def create_task(request: AlgorithmRequest) -> AlgorithmMiddleResponse:
    """
    创建新的MVHAND算法任务并存储在Redis中
    
    Args:
        request (AlgorithmRequest): 包含任务ID、回调URL和输入参数的请求对象
        
    Returns:
        AlgorithmMiddleResponse: 包含任务初始状态的响应对象
        
    Raises:
        HTTPException: 当任务创建失败时抛出500错误
    """
    try:
        task_response = AlgorithmMiddleResponse(
            task_id=request.task_id,
            task_callback_url=request.task_callback_url,
            task_status=TaskStatus.PENDING,
            task_progress=0,
            input_params=request.input_params,
            metrics=[]
        )
        await redis_client.set_data(request.task_id, task_response.model_dump())
        return task_response
    except Exception as e:
        logger.error(f"创建MVHAND任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"创建MVHAND任务失败: {str(e)}")

async def add_task_to_background(
    background_tasks: BackgroundTasks,
    input_params: Dict[str, Any],
    task_id: str
) -> None:
    """
    将MVHAND任务添加到后台执行队列
    
    Args:
        background_tasks: FastAPI的后台任务对象
        input_params: 算法输入参数
        task_id: 任务唯一标识
    """
    background_tasks.add_task(run_task_in_background, input_params, task_id)

@router.post("/execute/{task_id}", response_model=AlgorithmMiddleResponse)
async def execute_task(
    task_id: str, 
    background_tasks: BackgroundTasks
) -> AlgorithmMiddleResponse:
    """
    执行指定ID的MVHAND任务
    
    Args:
        task_id: 任务唯一标识
        background_tasks: FastAPI的后台任务对象
        
    Returns:
        AlgorithmMiddleResponse: 任务执行状态的响应对象
        
    Raises:
        HTTPException: 当任务不存在时抛出404错误，执行失败时抛出500错误
    """
    try:
        task_data = await redis_client.get_data(task_id)
        if not task_data:
            raise HTTPException(status_code=404, detail="MVHAND任务不存在")

        input_params = task_data.get("input_params")
        result_key = f'{task_id}_result'
        
        # 更新任务状态为执行中
        running_status = {
            "task_id": task_id,
            "task_callback_url": task_data.get("task_callback_url"),
            "task_status": TaskStatus.RUNNING,
            "task_progress": 0,
            "input_params": input_params,
            "metrics": []
        }
        await redis_client.set_data(result_key, running_status)
        
        # 添加到后台任务队列
        await add_task_to_background(background_tasks, input_params, task_id)
        
        return AlgorithmMiddleResponse(**running_status)
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"执行MVHAND任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"执行MVHAND任务失败: {str(e)}")

@router.get("/result/{task_id}", response_model=Union[AlgorithmResponse, AlgorithmMiddleResponse])
async def get_result(task_id: str) -> Union[AlgorithmResponse, AlgorithmMiddleResponse]:
    """
    获取MVHAND任务执行结果
    
    Args:
        task_id: 任务唯一标识
        
    Returns:
        Union[AlgorithmResponse, AlgorithmMiddleResponse]: 
            - 如果任务完成，返回AlgorithmResponse
            - 如果任务进行中，返回AlgorithmMiddleResponse
            
    Raises:
        HTTPException: 当结果不存在时抛出404错误，获取失败时抛出500错误
    """
    try:
        result_key = f'{task_id}_result'
        result_data = await redis_client.get_data(result_key)
        if not result_data:
            raise HTTPException(status_code=404, detail="MVHAND结果不存在")
        
        return (AlgorithmResponse(**result_data) 
                if "output_params" in result_data 
                else AlgorithmMiddleResponse(**result_data))
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取MVHAND结果失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"获取MVHAND结果失败: {str(e)}")

async def run_task_in_background(input_params: Dict[str, Any], task_id: str) -> None:
    """
    在后台执行MVHAND算法任务
    
    Args:
        input_params: 算法输入参数
        task_id: 任务唯一标识
    """
    try:
        result = await service.run_mvhand(
            input_params=input_params,
            task_id=task_id,
        )
        await redis_client.update_data(f'{task_id}_result', result.model_dump())
    except Exception as e:
        logger.error(f"MVHAND后台任务执行失败: {str(e)}")
        # 更新任务状态为失败
        error_status = {
            "task_id": task_id,
            "task_status": TaskStatus.FAILED,
            "error_message": str(e)
        }
        await redis_client.update_data(f'{task_id}_result', error_status) 